from convdiffsolverbase import ConvDiffSolverBase
from dolfin import *

class Supg(ConvDiffSolverBase):
    '''Original SUPG method by Hughes&Brooks(1982).'''       
    def __init__(self):
        self.name = 'Supg'
        ConvDiffSolverBase.__init__(self)

    def solve(self,problem,get_tau):
        '''Function get_tau obtains stabilization parameter.'''
        problemName, mesh, eps, BCValues, b, f, BCIndicators = problem.get_vars()

        V = FunctionSpace(mesh,'CG',1) 
        VV = VectorFunctionSpace(mesh,'CG',1)
        
        DG = FunctionSpace(mesh,'DG',0)

        u = TrialFunction(V)
        varphi = TestFunction(V)

        v = interpolate(b,VV)

        tau = Function(DG)                           # stabilization parameters
        tau.vector()[:] = get_tau(v,eps)

        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                          ) # a way to realize swith on |v| == 0.
                        
        a = eps*inner(nabla_grad(u),nabla_grad(varphi))*dx + inner(v,nabla_grad(u))*varphi*dx +\
            tau*pert*inner(dot(nabla_grad(u),outer(v,v)),nabla_grad(varphi))*dx 
        
        l = f*varphi*dx + tau*pert*f*inner(v,nabla_grad(varphi))*dx

        bcs = [DirichletBC(V,value,where) for value, where in zip(BCValues, BCIndicators)]

        A,L = assemble_system(a,l,bcs)

        u = Function(V)
        solve(A,u.vector(),L)
        
        plot(u,interactive=True)
        self.save(u,problem)

